#include <mm/kheap.h>
#include <mm/paging.h>
#include <mm/kmalloc.h>
#include <utils/debug.h>

#define FIND_FOOTER(h)  (struct kheap_footer*) ((uint32)(h) + (h)->size - sizeof(struct kheap_footer))

extern struct page_directory *kernel_directory;

int32 find_hole(struct kheap* heap, uint32 size, uint8 page_align) {
    assert(heap);
    
    uint32 idx = 0;
    while (idx < heap->index.count) {
        struct kheap_header* header = (struct kheap_header*) sorted_array_get(&(heap->index), idx);
        if (page_align) {
            uint32 location = (uint32)header;
            location += sizeof(struct kheap_header);
            int32 offset = 0;
            if (IS_PAGE_ALIGNED(location)) {
                offset = 0x1000 - location % 0x1000;
            }
            int32 hole_size = header->size - offset;
            if (hole_size >= (int32)size) {
                return idx;
            }
        } else if (header->size >= size) {
            return idx;
        }
        idx ++;
    }
    
    return -1;
}

int8 header_compare(item a, item b) {
    struct kheap_header *ha = (struct kheap_header*)a;
    struct kheap_header *hb = (struct kheap_header*)b;
    
    return ha->size - hb->size;
}

struct kheap* kheap_create(uint32 start, uint32 end, uint32 max, uint8 flags) {
    struct kheap* heap = (struct kheap*) kmalloc(sizeof(struct kheap));
    
    assert (start % 0x1000 == 0 && "Start address is not page aligned");
    assert (end % 0x1000 == 0 && "End address is not page aligned");
    
    heap->index = sorted_array_create_in((item*)start, HEAP_INDEX_SIZE, &header_compare);
    
    start += sizeof(item) * HEAP_INDEX_SIZE;
    
    ENSURE_PAGE_ALIGNED(start);
    
    heap->start = start;
    heap->end = end;
    heap->max = max;
    heap->flags = flags;
    
    struct kheap_header *hole = (struct kheap_header*) start;
    hole->size = end - start;
    hole->magic = HEAP_MAGIC;
    hole->is_hole = 1;
    sorted_array_insert(&(heap->index), (item)hole);
    
    return heap;
}

uint32 expand(struct kheap* heap, uint32 new_size) {
    assert(heap);
    // we can't expand more than what we can hold
    assert(new_size > heap->end - heap->start);
    
    ENSURE_PAGE_ALIGNED(new_size);
    
    assert(heap->start + new_size <= heap->max);
    
    uint32 i = heap->end - heap->start; // start at the old size
    while (i < new_size) {
        struct page* p = get_page(kernel_directory, heap->start + i);
        if (!p) {
            p = new_page(kernel_directory, heap->start + i, 0x7);
        }
        alloc_frame(p, heap->flags);
        i += PAGE_SIZE;
    }
    
    heap->end = heap->start + new_size;
    return new_size;
}

uint32 contract(struct kheap* heap, uint32 new_size) {
    assert(heap);
    assert(new_size < heap->end - heap->start);
    
    ENSURE_PAGE_ALIGNED(new_size);
    
    if (new_size < HEAP_MIN_SIZE) {
        new_size = HEAP_MIN_SIZE;
    }
    uint32 i = heap->end - heap->start - PAGE_SIZE;
    while (i >= new_size) {
        free_frame(get_page(kernel_directory, heap->start + i));
        i -= PAGE_SIZE;
    }
    
    heap->end = heap->start + new_size;
    return new_size;
}

uint32 find_endmost_index(struct sorted_array* array) {
    uint32 i;
    uint32 max_idx = -1;
    item max_value = 0;
    for (i = 0; i < array->count; i++) {
        item tmp = sorted_array_get(array, i);
        if (tmp > max_value) {
            max_value = tmp;
            max_idx = i;
        }
    }
    
    return max_idx;
}

void* kheap_alloc_in_hole(struct kheap* heap, uint32 size, uint8 page_align, int32 hole_idx) {
    struct kheap_header* orig_hole_header = (struct kheap_header*) sorted_array_get(&heap->index, hole_idx);
    uint32 new_size = size + sizeof(struct kheap_header) + sizeof(struct kheap_footer);
    uint32 orig_hole_pos = (uint32)orig_hole_header;
    uint32 orig_hole_size = orig_hole_header->size;
    
    if (orig_hole_size-new_size < sizeof(struct kheap_header)+sizeof(struct kheap_footer)) {
        size += orig_hole_size-new_size;
        new_size = orig_hole_size;
    }

    if (page_align && !IS_PAGE_ALIGNED(orig_hole_pos)) {
        uint32 new_location   = orig_hole_pos + PAGE_SIZE - (orig_hole_pos & 0xFFF) - sizeof(struct kheap_header);
        // setup header
        struct kheap_header *hole_header = (struct kheap_header *)orig_hole_pos;
        hole_header->size = PAGE_SIZE - (orig_hole_pos & 0xFFF) - sizeof(struct kheap_header);
        hole_header->magic = HEAP_MAGIC;
        hole_header->is_hole = 1;
        // setup footer
        struct kheap_footer *hole_footer = FIND_FOOTER(hole_header);
        hole_footer->magic = HEAP_MAGIC;
        hole_footer->header = hole_header;
        orig_hole_pos = new_location;
        orig_hole_size = orig_hole_size - hole_header->size;
    } else {
        sorted_array_remove(&heap->index, hole_idx);
    }

    struct kheap_header *block_header  = (struct kheap_header *)orig_hole_pos;
    block_header->magic     = HEAP_MAGIC;
    block_header->is_hole   = 0;
    block_header->size      = new_size;

    struct kheap_footer *block_footer  = FIND_FOOTER(block_header);
    block_footer->magic      = HEAP_MAGIC;
    block_footer->header    = block_header;

    // let's split, only if remaining size is positive
    if (orig_hole_size > new_size) {
        struct kheap_header *hole_header = (struct kheap_header *) (orig_hole_pos + sizeof(struct kheap_header) + size + sizeof(struct kheap_footer));
        hole_header->magic = HEAP_MAGIC;
        hole_header->is_hole = 1;
        hole_header->size     = orig_hole_size - new_size;
        struct kheap_footer *hole_footer = FIND_FOOTER(hole_header);
        if ((uint32)hole_footer < heap->end) {
            hole_footer->magic = HEAP_MAGIC;
            hole_footer->header = hole_header;
        }
        sorted_array_insert(&heap->index, (item)hole_header);
    }
    
    return (void *) ((uint32)block_header+sizeof(struct kheap_header));
}

void *kheap_alloc_expanding(struct kheap *heap, uint32 size, uint8 page_align) {
    uint32 new_size = size + sizeof(struct kheap_header) + sizeof(struct kheap_footer);
    
    // Save some previous data.
    uint32 old_length = heap->end - heap->start;
    uint32 old_end = heap->end;

    uint32 new_length = expand(heap, old_length + new_size);
    uint32 max_idx = find_endmost_index(&heap->index);
    
    // didn't found any headers
    if (max_idx == (uint32)-1) {
        struct kheap_header *header = (struct kheap_header*) old_end;
        header->magic = HEAP_MAGIC;
        header->is_hole = 1;
        header->size = new_length - old_length;
        
        struct kheap_footer* footer = FIND_FOOTER(header);
        footer->magic = HEAP_MAGIC;
        footer->header = header;
        sorted_array_insert(&heap->index, (item) header);
    } else {
        // The last header needs adjusting.
        struct kheap_header *header = (struct kheap_header *) sorted_array_get(&heap->index, max_idx);
        header->size += new_length - old_length;
        // Rewrite the footer
        struct kheap_footer *footer = FIND_FOOTER(header);
        footer->header = header;
        footer->magic = HEAP_MAGIC;
    }

    // We now have enough space at the end of the index.
    return kheap_alloc_in_hole(heap, size, page_align, heap->index.count - 1);
}

void* kheap_alloc(struct kheap* heap, uint32 size, uint8 page_align) {
    assert(heap);
    
    
    uint32 new_size = size + sizeof(struct kheap_header) + sizeof(struct kheap_footer);
    
    int32 idx = find_hole(heap, new_size, page_align);
    if (idx == -1) {
        // no hole found, let's expand
        return kheap_alloc_expanding(heap, size, page_align);
    }
    
    return kheap_alloc_in_hole(heap, size, page_align, idx);
}

void kheap_free(struct kheap *heap, void *p) {
    if (p == 0) {
        return;
    }
    
    struct kheap_header *header = (struct kheap_header*) ((uint32)p - sizeof(struct kheap_header));
    struct kheap_footer *footer = FIND_FOOTER(header);

    // Sanity checks.
    assert(header->magic == HEAP_MAGIC);
    assert(footer->magic == HEAP_MAGIC);

    // Make us a hole.
    header->is_hole = 1;

    uint8 do_add = 1;

    // Unify left
    struct kheap_footer *test_footer = (struct kheap_footer*) ((uint32)header - sizeof(struct kheap_footer));
    if (test_footer->magic == HEAP_MAGIC && test_footer->header->is_hole == 1) {
        uint32 old_size = header->size;
        header = test_footer->header; // overwrite header with left footer's header
        footer->header = header;
        header->size += old_size;     // update size to include old header
        do_add = 0;
    }

    // Unify right
    // If the thing immediately to the right of us is a header...
    struct kheap_header *test_header = (struct kheap_header*) ((uint32)footer + sizeof(struct kheap_footer));
    if (test_header->magic == HEAP_MAGIC && test_header->is_hole) {
        header->size += test_header->size; // Increase our size.
        test_footer = FIND_FOOTER(test_header);
        footer = test_footer;
        
        // Find and remove this header from the index.
        uint32 idx = sorted_array_find(&heap->index, (item)test_header);
        if (idx != (uint32)-1){
            sorted_array_remove(&heap->index, idx);
        }
    }

    // If the footer location is the end address, we can contract.
    if ( (uint32)footer + sizeof(struct kheap_footer) == heap->end) {
        uint32 old_length = heap->end - heap->start;
        uint32 new_length = contract(heap, (uint32)header - heap->start);
        // Check how big we will be after resizing.
        if (header->size - (old_length - new_length) > 0) {
            // We will still exist, so resize us.
            header->size -= old_length-new_length;
            footer = (struct kheap_footer*) ( (uint32)header + header->size - sizeof(struct kheap_footer) );
            footer->magic = HEAP_MAGIC;
            footer->header = header;
        } else {
            // We will no longer exist. Remove us from the index.
            uint32 idx = sorted_array_find(&heap->index, test_header);
            if (idx != (uint32)-1) {
                sorted_array_remove(&heap->index, idx);
            }
        }
    }

    // If required, add us to the index.
    if (do_add == 1) {
        sorted_array_insert(&heap->index, header);
    }
}


void kheap_print(struct kheap *heap) {
#ifdef DEBUG_MALLOC
    kprintf("Start: %x\n", heap->start);
    kprintf("End: %x\n", heap->end);
    kprintf("Max: %x\n", heap->max);
    kprintf("----Index----\n");
    int idx = 0;
    while (idx < heap->index.count) {
        struct kheap_header *header = (struct kheap_header*)heap->index.array[idx];
        struct kheap_footer *footer = FIND_FOOTER(header);
        kprintf("  Header at %x\n", header);
        kprintf("    Addr: %x\n", ((uint32)header + sizeof(struct kheap_header)));
        kprintf("    Magic: %x\n", header->magic);
        kprintf("    Is Hole: %d\n", header->is_hole);
        kprintf("    Size: %d bytes\n", header->size);
        kprintf("  Footer at %x\n", footer);
        kprintf("    Header at %x\n", footer->header);
        kprintf("    Magic: %x\n", header->magic);
        kprintf("---------------\n");
        idx++;
    }
#endif
}


